package org.cupk2.service.utils;

import com.cupk2.llm.client.inter.Cupk2EmbedClientInter;
import com.cupk2.llm.client.inter.Cupk2LlmClientInter;
import com.cupk2.lucene.Cupk2LuceneFace;
import com.cupk2.lucene.query.DocumentQuery;
import com.cupk2.lucene.query.rule.DocumentRuleByVector;
import com.cupk2.lucene.result.DocumentResult;
import com.cupk2.lucene.utils.CupkLuceneUtils;
import org.cupk2.service.lucene.doc.enums.DocumentTitleEnum;
import org.cupk2.tools.db.ResultDataHandle;
import org.cupk2.tools.web.dto.QuotingKnowDto;
import org.cupk2.tools.web.dto.SendMsgDto;
import org.cupk2.tools.caches.CupkCacheKeys;
import org.cupk2.tools.caches.CupkCaches;
import org.cupk2.tools.db.commons.CupkUUID;

import java.util.Map;

/**
 * 缓存对话消息
 */
public class CupkAIMsgCacheUtils {

    public static String cacheMsg(SendMsgDto msg) {
        String talkKey = CupkUUID.getUUID32();
        msg.setTalkKey(talkKey);
        CupkCaches.getInstance().putCacheData(talkKey, msg, CupkCacheKeys.TALK_MSG);
        return talkKey;
    }

    public static SendMsgDto getMsg(String talkKey, Cupk2LlmClientInter talkClient, Cupk2EmbedClientInter embedClient) {
        SendMsgDto msg = CupkCaches.getInstance().getCacheData(talkKey, CupkCacheKeys.TALK_MSG);
        if ("all".equals(msg.getSearchKnow())) {

            DocumentQuery query = new DocumentQuery();
            query.addRule(new DocumentRuleByVector(DocumentTitleEnum.embedding_vector.name(), embedClient.getEmbedding(msg.getMsg()), 8));
            DocumentResult result = Cupk2LuceneFace.getSearchService(CupkLuceneUtils.getLuceneEmbeddingFilePath()).search(query);
            result.runDataHandle(new ResultDataHandle() {
                @Override
                public void handle(Map<String, Object> row) {
                    String text = (String) row.get(DocumentTitleEnum.text.name());
                    String knowid = (String) row.get(DocumentTitleEnum.knowId.name());
                    String title = (String) row.get(DocumentTitleEnum.title.name());
                    String fileid = (String) row.get(DocumentTitleEnum.fileId.name());
                    msg.getQuotingKnows().add(new QuotingKnowDto(knowid, fileid, title, text));
                }
            });
        }
        for (String key : msg.getHisTalkKeys()) {
            SendMsgDto hismsg = CupkCaches.getInstance().getCacheData(key, CupkCacheKeys.TALK_MSG);
            msg.getHisTalks().add(hismsg);
        }
        return msg;
    }
}
